import re
import time
from get_sample import get_func_sample
prompt = '''
我会给你五个函数及其对应的市场，说明，以及一个问题。你的任务是分析这个问题是否可以通过调用这些函数来解决。

具体要求如下：
1. 如果函数的功能可以完全解决问题，即函数的输出能够满足或超过问题的需求，你需要调用该函数。
2. 如果没有任何函数能够完全解决问题，你需要输出“无匹配函数”。

调用函数的方式如下：
```python
arg1 = xxx
arg2 = xxx
func(arg1, arg2)
```

示例：
```python
test_code = 71598
test_year = 2019
query_employee_total(test_code, test_year)
```

请按照以下格式进行分析和输出：
找到可以完全解决问题的函数时：
```python
代码
```
无合适函数时:
`无匹配函数`
'''
from kernel import CodeKernel, run_code
from llm import llm

def FuncSolver(question,sample):


    function = sample.get_sample(question,5)
    function_str =  '\n\n'.join(function)
    messages = [{'role':'system','content':prompt},
                {'role':'user','content':f"""
可选函数为：
```python
{function_str}
```
问题为：{question}
禁止做参数假设，缺少参数视为`无匹配函数`
请判断并输出:
"""},]

    answer = llm(messages)
    if '```python' in answer and '无匹配函数' not in answer:

        ckernel = CodeKernel()
        ckernel.start()
        while True:
            time.sleep(1)
            tr = run_code('''print("test output")''', ckernel, True)
            # print(f'开始测试输出：```{tr}```')
            if tr == 'test output':
                break
        run_code('''
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "last_expr"
# 简化异常显示
%xmode Minimal
''', ckernel, True)
        run_code('''
import warnings
warnings.filterwarnings("ignore")
import pandas as pd
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 20)
pd.set_option('display.float_format', lambda x: str(x))
from Models import *
from sqlalchemy import *
from utils.utils import exec_orm_sql
engine = create_engine('mysql+pymysql://')
''', ckernel, True)

        # 注册函数
        for i in function:
            code = i.split('if __name__ == "__main__":')[0]
            run_code(code, ckernel, True)


        code = re.search('```python(.*?)```', answer,re.DOTALL).group(1)

        res = run_code(code, ckernel, True)
        if not re.search('Empty DataFrame|Tracebcak|Error',res):
            return res


if __name__ == '__main__':
    sample = get_func_sample()
    FuncSolver()